import { Quat, Texture2D, Vec3, Vec4 } from "cc";
import { TensorProxy } from "../Inference/TensorProxy";
import { TensorShape } from "../../Core/Tensor";
import { ObservationSpec } from "./ObservationSpec";
import { Texture2DWrapper } from "../Utils/Texture2DWrapper";


/**
 * @en
 *  Allows sensors to write to both TensorProxy and float arrays/lists.
 * 
 * @zh
 * 允许传感器同时写入 TensorProxy 和浮点数组/列表。
 */
export class ObservationWriter {

    private _data: Array<number> | null = null!;
    private _offset: number = 0;

    private _proxy: TensorProxy | null = null!;
    private _batch: number = 0;

    private _tensorShape: TensorShape = null!;

    setTarget(data: TensorProxy, shape: number, offset: number): void;
    setTarget(data: number[], shape: number, offset: number): void;
    setTarget(data: number[], shape: number[], offset: number): void;
    setTarget(data: number[], shape: ObservationSpec, offset: number): void;
    setTarget(data: number[] | TensorProxy, shape: ObservationSpec | number[] | number, offset: number) {
        if (shape instanceof ObservationSpec) {
            const shapeArray = shape.shape;
            if (shapeArray.length === 1) {
                this._tensorShape = TensorShape.createWithArray([this._batch, shapeArray[0]]);
            } else if (shapeArray.length === 2) {
                this._tensorShape = TensorShape.createWithArray([this._batch, 1, shapeArray[0], shapeArray[1]]);
            } else {
                this._tensorShape = TensorShape.createWithArray([this._batch, shapeArray[0], shapeArray[1], shapeArray[2]]);
            }

            this._data = data as number[];
            this._offset = offset;
            this._proxy = null;
        } else if (typeof shape === "number") {
            this._proxy = data as TensorProxy;
            this._batch = shape;
            this._data = null;
            this._tensorShape = this._proxy.data.shape;
        } else if (shape instanceof Array) {

            this._data = data as number[];
            this._offset = offset;
            this._proxy = null;
            this._batch = 0;


            if (shape.length === 1) {
                this._tensorShape = TensorShape.createWithArray([this._batch, shape[0]]);
            } else if (shape.length === 2) {
                this._tensorShape = TensorShape.createWithArray([this._batch, 1, shape[0], shape[1]]);
            } else {
                this._tensorShape = TensorShape.createWithArray([this._batch, shape[0], shape[1], shape[2]]);
            }


        }

    }

    /**
     * 1D write access at a specified index. Use addList if possible instead.
     * @param index Index to write to.
     */
    public setItem(index: number, value: number): void {
        if (this._data !== null) {
            this._data[index + this._offset] = value;
        } else {
            this._proxy?.data.setCache([this._batch, index + this._offset], value);
        }
    }

    /**
     * 3D write access at the specified height, width, and channel.
     * @param h Height.
     * @param w Width.
     * @param ch Channel.
     */
    public setItem3D(h: number, w: number, ch: number, value: number): void {
        if (this._data !== null) {
            if (h < 0 || h >= this._tensorShape.height) {
                throw new RangeError(`height value ${h} must be in range [0, ${this._tensorShape.height - 1}]`);
            }
            if (w < 0 || w >= this._tensorShape.width) {
                throw new RangeError(`width value ${w} must be in range [0, ${this._tensorShape.width - 1}]`);
            }
            if (ch < 0 || ch >= this._tensorShape.channels) {
                throw new RangeError(`channel value ${ch} must be in range [0, ${this._tensorShape.channels - 1}]`);
            }

            const index = this._tensorShape.index4D(this._batch, h, w, ch + this._offset);
            this._data[index] = value;
        } else {
            this._proxy?.data.setCache([this._batch, h, w, ch + this._offset], value);
        }
    }

    /**
     * Write the list of floats.
     * @param data The actual list of floats to write.
     * @param writeOffset Optional write offset to start writing from.
     */
    public addList(data: Array<number>, writeOffset: number = 0): void {
        if (this._data !== null) {
            for (let index = 0; index < data.length; index++) {
                const val = data[index];
                this._data[index + this._offset + writeOffset] = val;
            }
        } else {
            for (let index = 0; index < data.length; index++) {
                const val = data[index];
                this._proxy?.data.setCache([this._batch, index + this._offset + writeOffset], val);
            }
        }
    }

    /**
     * Write the Vector3 components.
     * @param vec The Vector3 to be written.
     * @param writeOffset Optional write offset.
     */
    public addVector3(vec: Vec3, writeOffset: number = 0): void {
        if (this._data !== null) {
            this._data[this._offset + writeOffset + 0] = vec.x;
            this._data[this._offset + writeOffset + 1] = vec.y;
            this._data[this._offset + writeOffset + 2] = vec.z;
        } else {
            this._proxy?.data.setCache([this._batch, this._offset + writeOffset + 0], vec.x);
            this._proxy?.data.setCache([this._batch, this._offset + writeOffset + 1], vec.y);
            this._proxy?.data.setCache([this._batch, this._offset + writeOffset + 2], vec.z);
        }
    }

    /**
     * Write the Vector4 components.
     * @param vec The Vector4 to be written.
     * @param writeOffset Optional write offset.
     */
    public addVector4(vec: Vec4, writeOffset: number = 0): void {
        if (this._data !== null) {
            this._data[this._offset + writeOffset + 0] = vec.x;
            this._data[this._offset + writeOffset + 1] = vec.y;
            this._data[this._offset + writeOffset + 2] = vec.z;
            this._data[this._offset + writeOffset + 3] = vec.w;
        } else {
            this._proxy?.data.setCache([this._batch, this._offset + writeOffset + 0], vec.x);
            this._proxy?.data.setCache([this._batch, this._offset + writeOffset + 1], vec.y);
            this._proxy?.data.setCache([this._batch, this._offset + writeOffset + 2], vec.z);
            this._proxy?.data.setCache([this._batch, this._offset + writeOffset + 3], vec.w);
        }
    }

    /**
     * Write the Quaternion components.
     * @param quat The Quaternion to be written.
     * @param writeOffset Optional write offset.
     */
    public addQuaternion(quat: Quat, writeOffset: number = 0): void {
        if (this._data !== null) {
            this._data[this._offset + writeOffset + 0] = quat.x;
            this._data[this._offset + writeOffset + 1] = quat.y;
            this._data[this._offset + writeOffset + 2] = quat.z;
            this._data[this._offset + writeOffset + 3] = quat.w;
        } else {
            this._proxy?.data.setCache([this._batch, this._offset + writeOffset + 0], quat.x);
            this._proxy?.data.setCache([this._batch, this._offset + writeOffset + 1], quat.y);
            this._proxy?.data.setCache([this._batch, this._offset + writeOffset + 2], quat.z);
            this._proxy?.data.setCache([this._batch, this._offset + writeOffset + 3], quat.w);
        }
    }
}

/**
 * Provides extension methods for the ObservationWriter.
 */
export class ObservationWriterExtension {
    /**
     * Writes a Texture2D into an ObservationWriter.
     * @param obsWriter Writer to fill with Texture data.
     * @param texture The texture to be put into the tensor.
     * @param grayScale If set to true, the textures will be converted to grayscale before being stored in the tensor.
     * @returns The number of floats written.
     */
    public static writeTexture(obsWriter: ObservationWriter, texture: Texture2D, grayScale: boolean): number {
        if (texture.getPixelFormat() === Texture2D.PixelFormat.RGB888) {
            return ObservationWriterExtension.writeTextureRGB24(obsWriter, texture, grayScale);
        }
        const width = texture.width;
        const height = texture.height;
        const texturePixels = Texture2DWrapper.getPixels32(texture);
        // During training, we convert from Texture to PNG before sending to the trainer, which has the
        // effect of flipping the image. We need another flip here at inference time to match this.
        for (let h = height - 1; h >= 0; h--) {
            for (let w = 0; w < width; w++) {
                const currentPixel = texturePixels[(height - h - 1) * width + w];
                if (grayScale) {
                    obsWriter.setItem3D(h, w, 0, (currentPixel.r + currentPixel.g + currentPixel.b) / 3 / 255);
                } else {
                    // For Color32, the r, g and b values are between 0 and 255.
                    obsWriter.setItem3D(h, w, 0, currentPixel.r / 255);
                    obsWriter.setItem3D(h, w, 1, currentPixel.g / 255);
                    obsWriter.setItem3D(h, w, 2, currentPixel.b / 255);
                }
            }
        }
        return height * width * (grayScale ? 1 : 3);

    }

    private static writeTextureRGB24(obsWriter: ObservationWriter, texture: Texture2D, grayScale: boolean): number {
        const width = texture.width;
        const height = texture.height;
        const rawBytes = texture.image?._nativeAsset;
        // During training, we convert from Texture to PNG before sending to the trainer, which has the
        // effect of flipping the image. We need another flip here at inference time to match this.
        for (let h = height - 1; h >= 0; h--) {
            for (let w = 0; w < width; w++) {
                const offset = (height - h - 1) * width + w;
                const r = rawBytes[3 * offset];
                const g = rawBytes[3 * offset + 1];
                const b = rawBytes[3 * offset + 2];
                if (grayScale) {
                    obsWriter.setItem3D(h, w, 0, (r + g + b) / 3 / 255);
                } else {
                    // For Color32, the r, g and b values are between 0 and 255.
                    obsWriter.setItem3D(h, w, 0, r / 255);
                    obsWriter.setItem3D(h, w, 1, g / 255);
                    obsWriter.setItem3D(h, w, 2, b / 255);
                }
            }
        }
        return height * width * (grayScale ? 1 : 3);
    }
}



